/*
 * Code for blog.techhead.biz
 * Distributed under BSD-style license
 */

package biz.techhead.funcy;

/**
 *
 * @author Jonathan Hawkes <jhawkes at techhead.biz>
 */
public class Restarts {
    private Restarts() {}
            
    public static abstract class RestartableFuncE<O,I,E extends Throwable>
            implements FuncE<O,I,E> {

        public RestartableFuncE<O,I,E> restartCase(final Object restartKey,
                final Func<? extends O,Object> altFunc) {
            return bindRestart(restartKey, altFunc);
        }

        public RestartableFuncE<O,I,E> restartCase(final Object restartKey,
                final FuncE<? extends O,Object,? extends E> altFunc) {
            return bindRestart(restartKey, altFunc);
        }

        private RestartableFuncE<O,I,E> bindRestart(final Object restartKey,
                final FuncE<? extends O,Object,? extends Throwable> altFunc) {
            final FuncE<O,I,E> code = this;
            return new RestartableFuncE<O,I,E>() {
                @Override public O call(I in) throws E {
                    ThreadContext ctx = threadCtx.get();
                    try {
                        ctx.pushRestart(restartKey, altFunc);
                        return code.call(in);
                    } catch (StackUnwinder u) {
                        u.unwindStack();
                        return ctx.<O,E>restart();
                    } catch (Throwable e) {
                        return ctx.<O,E>tryToRestartFrom(e);
                    } finally {
                        ctx.popRestart();
                    }
                }
            };
        }

        public <T extends Throwable> RestartableFuncE<O,I,E>
                handlerBind(final Class<T> exType,
                final FuncE<?,? super T,RuntimeException> exHandler) {
            final FuncE<O,I,E> code = this;
            return new RestartableFuncE<O,I,E>() {
                @Override public O call(I in) throws E {
                    ThreadContext ctx = threadCtx.get();
                    try {
                        ctx.pushHandler(exType, exHandler);
                        return code.call(in);
                    } finally {
                        ctx.popHandler();
                    }
                }
            };
        }
    }

    @SuppressWarnings("unchecked")
    public static <O,I,E extends Throwable> RestartableFuncE<O,I,E>
            on(final FuncE<? extends O,? super I,? extends E> code) {
        if (code instanceof RestartableFuncE)
            return ( RestartableFuncE<O,I,E> ) code;
        return new RestartableFuncE<O,I,E>() {
            @Override public O call(I in) throws E {
                return code.call(in);
            }
        };
    }
    
    public static boolean invokeLater(Object restartKey) {
        return invokeLater(restartKey, Tuples.T0);
    }
    
    public static boolean invokeLater(Object restartKey, Object arg) {
        ThreadContext ctx = threadCtx.get();
        return ctx.invokeLater(restartKey, arg);
    }
    
    private static class ThreadContext {
        
        private Stack< Handler<?> > handlers;
        private Stack< RestartsScope > restarts;

        private Handler<?> handler;
        private FuncE<?,Object,? extends Throwable> restart;
        private Object restartArg;
        private int restartPoint;
        private int currentRestartPoint = 0;
        
        private Throwable ignore;

        ThreadContext() {
            restarts = new Stack<RestartsScope>();
            restarts.value = new RestartsScope();
        }
        
        boolean invokeLater(Object restartKey, Object arg) {
            if ( handler.findRestart(restartKey) ) {
                restartArg = arg;
                return true;
            }
            return false;
        }

        @SuppressWarnings("unchecked")
        <T,E extends Throwable> T restart() throws E {
            // If there is an exception here, it will travel down the stack
            // in place of the original. This is expected behavior. Right?
            try {
                return (T) restart.call(restartArg);
            } catch (Throwable t) {
                this.<E>castAndThrow(t);
                throw new AssertionError(); // unreachable statement
            }
        }

        <T,E extends Throwable> T tryToRestartFrom(Throwable t) throws E {
            
            // If we've already called handlers on this exception
            // higher up the stack, don't do it again.
            if (ignore != null) {
                if ( (t == ignore) &&
                        (currentRestartPoint == restartPoint) ) {
                    restartPoint--;
                    // keep propagating the exception
                    this.<E>castAndThrow(t);
                } else {
                    ignore = null;
                }
            }

            // no restart invoked yet
            restart = null;

            // Iterate over all the handlers
            // starting from closest to the exception
            // then moving down the stack.
            Stack< Handler<?> > stack = handlers;
            while (stack != null) {
                handler = stack.value;
                // If the handler is applicable, give it a whirl.
                // If the handler invoked a restart, then we're done.
                // Otherwise, we'll keep calling handlers.
                if ( handler.handle(t) && (restart != null) ) {
                    // If necessary, unwind the stack down to the
                    // point of the selected restart.
                    if (restartPoint != currentRestartPoint) {
                        throw new StackUnwinder
                                (currentRestartPoint - restartPoint);
                    }
                    // Invoke the restart and return its value as
                    // a substitute for that of the original
                    // failed block
                    return this.<T,E>restart();
                }
                stack = stack.next;
            }

            // If no handler invoked a restart,
            // then re-throw the original exception.
            ignore = t;
            restartPoint = currentRestartPoint - 1;
            this.<E>castAndThrow(t);
            throw new AssertionError(); //unreachable statement
        }

        @SuppressWarnings("unchecked")
        private <E extends Throwable> void castAndThrow(Throwable t) throws E {
            if (t instanceof RuntimeException) throw (RuntimeException) t;
            if (t instanceof Error) throw (Error) t;
            throw (E) t;
        }

        <T extends Throwable> void pushHandler(Class<T> exType,
                FuncE<?,? super T,RuntimeException> exHandler) {
            handlers = new Stack< Handler<?> >(handlers);
            handlers.value = new Handler<T>(exType, exHandler, restarts.value);
        }

        void popHandler() {
            handlers = handlers.next;
        }

        void pushRestart(final Object restartKey,
                         final FuncE<?,Object,? extends Throwable> altFunc) {
            RestartsScope newTail = new RestartsScope() {
                final int id = currentRestartPoint++;
                @Override public boolean findRestart(Object key) {
                    if ( key.equals(restartKey) ) {
                        restart = altFunc;
                        restartPoint = id;
                        return true;
                    }
                    return super.findRestart(key);
                }
            };
            restarts.value.next = newTail;
            restarts = new Stack<RestartsScope>(restarts);
            restarts.value = newTail;
        }
        
        void popRestart() {
            restarts = restarts.next;
            restarts.value.next = null;
            currentRestartPoint--;
        }
    }

    private static class Handler<T extends Throwable> {

        private final Class<T> exType;
        private final FuncE<?,? super T,RuntimeException> exHandler;
        private final RestartsScope restartsBefore;

        Handler(Class<T> exType,
                FuncE<?,? super T,RuntimeException> exHandler,
                RestartsScope restartsBefore) {
            this.exType = exType;
            this.exHandler = exHandler;
            this.restartsBefore = restartsBefore;
        }

        @SuppressWarnings("unchecked")
        boolean handle(Throwable ex) {
            if ( exType.isAssignableFrom( ex.getClass() ) ) {
                try {
                    exHandler.call( (T) ex );
                } catch (Throwable he) {
                    // A handler should not affect the execution path
                    // except for by invoking restarts.
                    // If this handler throws an exception, deal with it in
                    // the same way that a Thread deals with exceptions
                    // that escape the run method.
                    Thread currentThread = Thread.currentThread();
                    currentThread.getUncaughtExceptionHandler()
                            .uncaughtException(currentThread, he);
                }
                return true;  // we called it, successful or no
            }
            return false;
        }

        boolean findRestart(Object key) {
            RestartsScope scope = restartsBefore.next;
            return (scope == null) ?
                false : scope.findRestart(key);
        }
    }

    /**
     * Represents a scope in which a restart is visible.
     */
    private static class RestartsScope {
        RestartsScope next;
        boolean findRestart(Object key) {
            return (next == null) ?
                false : next.findRestart(key);
        }
    }

    /**
     * A rudimentary stack in the form of a singly-linked list.
     * @param T the type of item the stack holds
     */
    private static class Stack<T> {
        T value;
        Stack<T> next;
        Stack() {}
        Stack(Stack<T> next) {
            this.next = next;
        }
    }
        
    /**
     * This class does not really represent an error.
     * However, since most Java programmers are taught from day 1
     * not to suppress subclasses of Error, this seemed the
     * best way to navigate down the stack undeterred.
     */
    private static class StackUnwinder extends Error {
        private int toUnwind;
        StackUnwinder(int toUnwind) {
            this.toUnwind = toUnwind;
        }
        void unwindStack() {
            if (--toUnwind != 0) throw this;
        }
    }

    private static ThreadLocal<ThreadContext> threadCtx =
            new ThreadLocal<ThreadContext>() {
        @Override public ThreadContext initialValue() {
            return new ThreadContext();
        }
    };
}
